open Filename
open Printf
(*open Next
open Partition
open Linear*) 
open Unix
open Circuit
open Str
open Ext

let network = ref ""
let modelfile = ref ""
let outfile = ref ""
let outfile_js = ref ""
let outfile_dot = ref ""
let usage = "Usage: spn2ac -m <model_directory> -o <ac_output>"
let args = Arg.align
  ([
    ("-m",  Arg.Set_string modelfile,  " Input SPN (SPAC) model directory");
    ("-o",  Arg.Set_string outfile,  " Output arithmetic circuit");
(*    ("-js",  Arg.Set_string outfile_js,  " Output javascript graph");
    *)
    ("-dot",  Arg.Set_string outfile_dot,  " Output dot graph")] 
    @ common_arguments)



exception HSplitNotPossible
exception VSplitNotPossible
exception NodeTypeError


type ndetails = TimesNode
        | PlusNode
        | NullNode 
  | ACNode;;


type spnode = {
  mutable id: int;
  mutable parents: int array;
  mutable children2: int array;
  mutable params: float array;
  mutable nodetype: ndetails;
  mutable schema: (int * int) array;
  mutable data: int array array;
  mutable acname: string;
  mutable ac_index: int ;
  mutable counts: int array;
  mutable activated_link: int;
  mutable kl: float;
  mutable thresh: float;
}

(*let root = ref spnode
*)


let node_array = ref [||]
let comps = ref [||] 
let comps2 = ref [||] 
let feature_lists = ref [||]
let scratches = ref [||] 
let logzs = ref [||] 
(*let data_parts = ref [||]
*)
let get_next_id () = 
  Array.length !node_array



let create_node t pt c pm d s=
  let n = {id = get_next_id(); 
  parents = pt;
  children2 = c; 
  params = pm;
  nodetype = t;
  schema = s; 
  data = d;
  acname = "";
  ac_index = 0;
  counts = [||];
  activated_link = 0;
  kl = 0.0;
  thresh = 0.0} in
  node_array := Array.append !node_array [|n|];
  n


let create_acnode = create_node ACNode [||] [||] [||]  

let create_nullnode = create_node NullNode [||] [||] [||]  
 
(*
let append_node node =

  ignore()
*)

let create_plus = create_node PlusNode [||] [||] [||] 


let create_times  = create_node TimesNode [||] [||] [||] 


let len a = Array.length a 

let node_as_array n = [|n|]

let node_as_list n = [n]

let add_child node child = 
  node.children2 <- Array.append node.children2  [|child.id|];
  node.params <- Array.append node.params [|log(float_of_int(Array.length child.data)) -. log( float_of_int(Array.length node.data))|];
  node.counts <- Array.append node.counts [|(Array.length child.data)|];
  ignore()

let set_parent node parent =
  add_child parent node;
  node.parents <- Array.append node.parents [|parent.id|];
  ignore()

let print_params node =
  Array.iter (fun p->printf "%f " (exp p)) node.params 

let rec print_tree node depth nt = 
  for i=0 to depth do
    printf "  "
  done;
  (*printf "id: %d sample: %d vars: %d type: " node.id (Array.length node.data) (Array.length node.schema);*)
  printf "id: %d vars: %d type: " node.id (*(Array.length node.data)*) (Array.length node.schema);
  (match node.nodetype with ACNode-> printf "ac\n" | TimesNode-> printf "* \n" (* "kl: %f, thr: %f\n" (node.kl/.nt) node.thresh*) | PlusNode-> printf "+ "; print_params node; printf "\n" (*; printf "kl: %f\n" (node.kl/.nt)*) | NullNode -> ignore());
  let n = len node.children2 in
  for i=0 to n - 1 do
    print_tree !node_array.(node.children2.(i)) (depth+1) nt
  done;
  ignore()


let print_spn node nt =
  print_tree node 0 nt



let minisleep (sec: float) =
    ignore (Unix.select [] [] [] sec)
exception CircuitLoadException

let get_node_type t = match t with ACNode->"ac" | PlusNode-> "+"| TimesNode ->"*" | NullNode->"null"

let load_circuit filename =
  try Circuit.load_with_features (open_in filename) 
  with Circuit.Unknown_child(n,c) -> printf "Error: node %s references unknown child node %s.\n" n c; raise CircuitLoadException

let load_comps () =
  feature_lists := [||];
  comps := [||];
  for i=1 to (Array.length !node_array) - 1  do
  (*  printf "node id:%d type:%s\n" !node_array.(i).id (match !node_array.(i).nodetype with ACNode->"ac" | PlusNode-> "+"| TimesNode ->"*");
  *)
    let nodetype = !node_array.(i).nodetype in
    if (nodetype == ACNode) then (* || nodetype == PlusNode || nodetype == TimesNode) then *)
    begin
      let comp_name = !node_array.(i).acname in
    (*  printf "%s\n" comp_name; flush_all();*)
    let full_comp_name = String.concat dir_sep [!modelfile; comp_name] in
      let (component, fl) = load_circuit full_comp_name in
      let n = Array.length !comps in 
      comps := Array.append !comps [|component|];
      feature_lists := Array.append !feature_lists [|fl|];
      !node_array.(i).ac_index <- n;
    end
  done;

  scratches := Array.map (Circuit.create_scratch) !comps;
  logzs := Array.map2 (fun s c->Circuit.compute_z s c ) !scratches !comps;
  ignore()




let get_type_name nodetype = 
  if nodetype == PlusNode then "+" else if nodetype == TimesNode then "*" else "ac"

exception InsufficientArguments


let print_node out node =
  let n = (Array.length node.children2) in
  fprintf out "%d %s\n" node.id  (match node.nodetype with ACNode-> "ac" | TimesNode-> "*" | PlusNode-> "+" | NullNode ->  "null");
  for i=0 to n-1 do
    fprintf out "%d " node.children2.(i)
  done;
  fprintf out "\n";
  for i=0 to n-1 do
    fprintf out "%f " node.params.(i)
  done;
  fprintf out "\n"

let print_model out root =
  fprintf out "%d\n" (Array.length !node_array);
  Array.iter (print_node out) !node_array

exception ParseError of string
let space = regexp " " 

let read_line_tokens model = split space (input_line model) 

let parse_node model = 
  let tokens = split space (input_line model) in
  (*printf "line: %s %s %s\n" (List.nth tokens 0) (List.nth tokens 1) (List.nth tokens 2); 
  *)
  let i = match (List.nth tokens 0) with "n" -> 0 | _-> raise (ParseError "not a node") in
  ignore(i);
  let nodeid = int_of_string (List.nth tokens 1) in
  let nodetype = match (List.nth tokens 2) with "+" -> PlusNode | "ac" -> ACNode | "*" -> TimesNode | "null"->NullNode | _-> raise(ParseError "node type incorrect") in
  let parentid = int_of_string (List.nth tokens 3) in
  (*printf "network: %s nodeid :%d nodetype:%s %s parentid :%d \n" !network nodeid (List.nth tokens 2) (get_node_type(nodetype)) parentid;
  *)
  let parents = if parentid < 0 then [||] else [|parentid|] in
  
  (match nodetype with PlusNode -> 
    begin
    let tokens = read_line_tokens model in
    let children2 = Array.of_list (List.map (int_of_string) tokens) in
    let tokens = read_line_tokens model in
    let params = Array.of_list (List.map (float_of_string) tokens) in
    
    let tokens = read_line_tokens model in
    let schema = Array.of_list (List.map (int_of_string) tokens ) in
    let i_schema = Array.map (fun s->(2,s))schema in
    let n = create_node  nodetype parents children2 params [||] i_schema in
    (*!node_array.(nodeid).acname <- sprintf "%s-%d.ac" !network nodeid;*)
    ignore(n)
    end
    | TimesNode ->
    begin

    let tokens = read_line_tokens model in
    let thresh = float_of_string(List.hd tokens) in
    let tokens = read_line_tokens model in
    let children2 = Array.of_list (List.map (int_of_string) tokens) in

    let tokens = read_line_tokens model in
    let schema = Array.of_list (List.map (int_of_string) tokens ) in
    let i_schema = Array.map (fun s->(2,s))schema in
    let n = create_node  nodetype parents children2 [||] [||] i_schema in
    n.thresh<-thresh;
    ignore(n)
    (*!node_array.(nodeid).acname <- sprintf "%s-%d.ac" !network nodeid;*)
    end
    | ACNode ->
    begin

    let tokens = read_line_tokens model in
    let schema = Array.of_list (List.map (int_of_string) tokens ) in
    let i_schema = Array.map (fun s->(2,s))schema in
    let ac_node = create_node nodetype parents [||] [||] [||] i_schema in
   
    ignore(ac_node);
    !node_array.(nodeid).acname <- sprintf "%s-%d.ac" !network nodeid;

    end
    | NullNode-> begin let nullnode=create_nullnode [||] [||] in ignore(nullnode) end)
 (* | _->raise (ParseError "Not a valid node")); *)
  
  (*!node_array.(nodeid).nodetype <- nodetype;*)
  (*!node_array.(nodeid).acname <- sprintf "%s-%d.ac" !network nodeid*)

let load_model modelfile =
  try
    let model = open_in modelfile in
    network := input_line model;
    let num_nodes = int_of_string (input_line model) in
    for i=1 to num_nodes do
      parse_node model;     
    done;
  with Failure _ -> printf "ERROR" 

let rec rebuild_ac node =
  (match node.nodetype 
    with ACNode-> begin
        let circ = !comps.(node.ac_index) in
        let logz = !logzs.(node.ac_index) in
        (* let coeff = -.logz in *)
        let coeff_nodes = ref [] in
        let remaining = ref logz in
        let log2 = log 2.0 in
        while !remaining > log2 do
          let node = Circuit.create_const (-.log2) in
          coeff_nodes := node::!coeff_nodes;
          remaining := !remaining -. log2;
        done;

        let node = Circuit.create_const (-. !remaining) in
        coeff_nodes := node::!coeff_nodes;
        (*printf "nid: %d logz:%f 1/z: %f\n" node.id logz (exp coeff);*)
        let children = circ.root::!coeff_nodes in
        
        (* OLD: let coeff_node = Circuit.create_const coeff in *)
        let new_root = Circuit.create_times children in
        new_root
     end
       | TimesNode-> begin
        let num = len node.children2 in
        let children_roots = ref [] in 
        for i=0 to num - 1 do
          let circ_root = rebuild_ac !node_array.(node.children2.(i)) in
          children_roots := circ_root::!children_roots
        done;
        Circuit.create_times (List.rev !children_roots)
     end
       | PlusNode-> begin
        let num = len node.children2 in
        let children_roots = ref [] in 
        for i=0 to num - 1 do
          let circ_root = rebuild_ac !node_array.(node.children2.(i)) in
          let param_node = Circuit.create_const node.params.(i) in
          let p_node = Circuit.create_times [param_node; circ_root] in
          children_roots := p_node::!children_roots
        done;
        Circuit.create_plus (List.rev !children_roots)
      end ) 

let of_graph4 s r =
  let vn = Circuit.make_vnodes s in
  let vnl = Array.map Array.to_list vn in
  (*let fixchild c =
    if is_var c then
      let (var, value) = var_details c in
      vn.(var).(value)
    else c in
  let relink_vars m =
    m.children <- List.map fixchild m.children in
  Circuit.node_iter relink_vars r ;
  *)
  Circuit.of_graph s vn vnl r
           

let rebuild_spn spnroot =
  let schema = Array.map fst spnroot.schema in
  
  let vn = Circuit.make_vnodes  schema in
  let vnl = Array.map Array.to_list vn in
    
  for i=0 to (Array.length !node_array)  -1 do
    if !node_array.(i).nodetype == ACNode then 
    begin
      let acnode = !node_array.(i) in
      let localmapping = Array.map snd acnode.schema in
      let circ = !comps.(acnode.ac_index) in
      let fixchild c =
        if is_var c then
          let (var, value) = var_details c in
          (*printf "n %d: %d->%d\n" i var localmapping.(var);
          *)
          vn.(localmapping.(var)).(value)
        else c in
      let relink_vars m =
        m.children <- List.map fixchild m.children in
      Circuit.node_iter relink_vars circ.root ;
      let newcirc = Circuit.of_graph (Array.map fst acnode.schema) vn vnl circ.root in
      (* let newindex = Array.length !comps2 in *)
      comps2 := Array.append !comps2 [|newcirc|];
      (* acnode.ac_index <- newindex; *)
      !comps.(acnode.ac_index) <- newcirc
    end
  done;
  ignore()



let main() =
  Arg.parse args ignore usage;
  common_log_init();

  if !modelfile = "" then (Arg.usage args usage; exit(-1));
  
  if String.suffix !modelfile dir_sep then
    modelfile := String.sub !modelfile 0 ((String.length !modelfile) - 1);
  
  if not (check_suffix !modelfile ".spn") then begin 
    printf "Only accepts .spn as the input model file"; exit 1 
  end;
  
  if !outfile_js <> "" && not (check_suffix !outfile_js ".js") then begin 
    printf "Javascript output file should end with .js\n"; 
    exit 1 
  end;

  load_model (String.concat dir_sep [!modelfile; "spac.m"]);  
  
  load_comps();
  

  rebuild_spn !node_array.(0);
  let newroot = rebuild_ac !node_array.(0) in
  
  let newCircuit = of_graph4 (Array.map fst !node_array.(0).schema)  newroot in
  
  
  if !outfile <> "" then begin
    let out_ch = open_out !outfile in
    Circuit.output out_ch newCircuit;
    close_out out_ch 
  end;

  
  if !outfile_js <> "" then begin
    let out_ch = open_out !outfile_js in
    Circuit.output_js out_ch newCircuit;
    close_out out_ch
  end;
 
  if !outfile_dot <> "" then begin
    let out_ch = open_out !outfile_dot in
    Circuit.output_dot out_ch newCircuit;
    close_out out_ch
  end;
 
  printf "Done!\n";
  ignore()

;;

main()
